#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: runyuanye@163.com
"""

import argparse
import os
import sys
import numpy as np
import cv2
from collections import OrderedDict
import time
import random

sys.path.append('.')

"""
    手工合并相同的人脸
    
    按‘g’或‘G’键将当前所有人脸设为同一FaceID，以便将不同人脸合并成同一FaceID，默认是都不合并
    按‘e’或‘E’键将当前人脸删除，即最终合并人脸FaceID时被删除不会输出，当不是人脸或质量不好时可选择删除，
    按‘t’或‘T’键将当前所有人脸删除，
    按‘d’或‘D’键切换下一页，
"""

h_image_count = 10
v_image_count = 10
img_width = 96*2
img_height = 96*2
# img_width = 96
# img_height = 96

batch_size = 32

ix, iy = -1, -1
draw_image = None
image_bboxes = []
select_image_bbox_idxes = set()
current_mouse_xy = None


def get_xy_bbox_idx(xy):
    global image_bboxes
    bbox_idx = None
    if xy is None:
        return bbox_idx
    x, y = xy
    for i, bbox in enumerate(image_bboxes):
        # print(bbox, x, y)
        if (bbox[0] < x < bbox[2]) and (bbox[1] < y < bbox[3]):
            bbox_idx = i
            break
    return bbox_idx


def mouse_event(event, x, y, flags, param):
    global draw_image, image_bboxes, select_image_bbox_idxes, current_mouse_xy

    current_mouse_xy = (x, y)

    # if draw_image is not None:
    #     if event == cv2.EVENT_LBUTTONDOWN:
    #         # print(x, y)
    #         bbox_idx = get_xy_bbox_idx(x, y)
    #         if bbox_idx is not None:
    #             if bbox_idx not in select_image_bbox_idxes:
    #                 select_image_bbox_idxes.add(bbox_idx)
    #             else:
    #                 select_image_bbox_idxes.remove(bbox_idx)



def image_resize(img, img_width, img_height, mode=1):
    if mode == 0:
        img = cv2.resize(img, (img_width, img_height))
    elif mode == 1:
        w, h = img.shape[1], img.shape[0]
        wscale = img_width / w
        hscale = img_height / h
        scale = min(wscale, hscale)
        w = int(round(w * scale))
        h = int(round(h * scale))
        image = cv2.resize(img, (w, h))
        img = np.full((img_height, img_width, image.shape[2]), fill_value=0, dtype=np.uint8)
        x = (img_width - w) // 2
        y = (img_height - h) // 2
        img[y:y + h, x:x + w, :] = image
    return img


def show_images(image_dir, file_names, dis_text='', win_name='image', wait=False):
    h_image_count = 20
    v_image_count = 10

    display_image_count = h_image_count * v_image_count
    files = []
    total_image_count = len(file_names)
    last_idx = total_image_count - 1
    # if total_image_count > 0:
    #     print('{}: {}'.format(win_name, file_names[0]))
    for idx, file_name in enumerate(file_names):
        files.append(file_name)
        file_count = len(files)
        if file_count < display_image_count and (idx < last_idx):
            continue
        imgs = []
        for file_name in files:
            file_path = os.path.join(image_dir, file_name.strip())
            img = cv2.imread(file_path)
            img = image_resize(img, img_width, img_height)
            imgs.append(img)
        files.clear()
        image_count = len(imgs)
        images = []
        x_off = 0
        y_off = 0
        for j, img in enumerate(imgs):
            if j % h_image_count == 0:
                x_off = 0
                if j != 0:
                    y_off += img_height
                h_images = [img]
            else:
                x_off += img_width
                h_images.append(img)
            if ((j + 1) % h_image_count == 0) or ((j + 1) == image_count):
                images.append(np.concatenate(h_images, axis=1))

        if len(images) > 1:
            w0 = images[0].shape[1]
            h0 = images[0].shape[0]
            w1 = images[-1].shape[1]
            h1 = images[-1].shape[0]
            if (w0 != w1) or (h0 != h1):
                # image = np.empty_like(images[0])
                image = np.zeros_like(images[0])
                image[0:h1, 0:w1, :] = images[-1]
                images[-1] = image
        draw_image = np.vstack(images)
        cv2.putText(draw_image, '%s %d/%d/%d' % (dis_text, idx + 1, last_idx + 1, total_image_count), (2, 12),
                    cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 255, 0), 1)
        cv2.imshow(win_name, draw_image)
        if wait:
            key = cv2.waitKey()
            if key == 27 or key == ord('q') or key == ord('Q'):
                return
        else:
            break


def read_file_list(file_list, mode=0):
    faceid_subfaceid_files = {}
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read File Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            lines = line.strip().split()

            if mode == 0:
                # 015ff272-3768-4d4c-bebc-c14dc93065cd 477fc604-1479-4745-b8e9-9440a6ac8187 f8e4735b44724e25077106e0eaef15f0.jpg 0
                faceid = lines[0]
                subfaceid = lines[1]
                file_name = lines[2]
            else:
                # 000016.jpg
                file_name = lines[0].strip()
                faceid = file_name.split('.')[0]
                subfaceid = faceid

            faceid_subfaceid = '{}_{}'.format(faceid, subfaceid)
            file_selects = faceid_subfaceid_files.get(faceid_subfaceid, None)
            if file_selects is None:
                faceid_subfaceid_files[faceid_subfaceid] = [file_name]
            else:
                file_selects.append(file_name)

        print('Read File Count: {}'.format(file_count))
    return faceid_subfaceid_files


def save_process(faceid_file_list_path, xiaoqu_faceid_file_nums, cur_id_idx_file_path, cur_id_idx):
    if len(xiaoqu_faceid_file_nums) > 0:
        with open(faceid_file_list_path, 'w') as file:
            for xiaoqu, faceids, files, nums in xiaoqu_faceid_file_nums:
                same_faceid_nums_info = xiaoqu + ':' + (','.join(faceids)) + ';' + (','.join(map(str, nums))) + '\n'
                file.write(same_faceid_nums_info)
        os.system('chmod a+wr {}'.format(faceid_file_list_path))

    with open(cur_id_idx_file_path, 'w') as file:
        process_info = '{}\n'.format(cur_id_idx)
        file.write(process_info)
    os.system('chmod a+wr {}'.format(cur_id_idx_file_path))


def read_merge_file(file_list, image_file_list, use_face_num=True):
    faceid_subfaceids_files = []
    file_count = 0
    images_file = open(image_file_list, 'r')
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read Feature Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            # NONE:c651b1e5-fdbc-4447-9a6a-9213d7729faa_ac638be4-0e9a-42e8-a29b-c730a41aac50,752524c4-e540-45e2-8388-55aab4d72150_94d61342-0841-432a-a8f2-5bb6cf5c9a4f
            lines = line.strip().split(':')
            xiaoqu = lines[0]
            if use_face_num:
                faceid_nums = lines[1].split(';')

                faceids = faceid_nums[0].split(',')

                facenums = faceid_nums[1].split(',')
                facenums = list(map(int, facenums))
            else:
                faceids = lines[1].split(',')
                facenums = [0] * len(faceids)

            # NONE:c326b4e9de7ea3853a0e1c1d4bb973d8.jpg,31f81112f544a5dc3614cdcd862d158e.jpg
            file_line = images_file.readline()
            file_lines = file_line.strip().split(':')
            files = file_lines[1].split(',')

            faceid_subfaceids_files.append([xiaoqu, faceids, files, facenums])
        print('Read Merge Count: {}'.format(file_count))
    images_file.close()
    return faceid_subfaceids_files


def main():
    global draw_image, image_bboxes, select_image_bbox_idxes, current_mouse_xy
    parser = argparse.ArgumentParser(description="Select Image")
    parser.add_argument(
        "--image_dir", default="/rootfs/media/kasim/Data/data/FaceImage/Image", help="path to image file", type=str
    )
    parser.add_argument(
        "--file_list", default="/rootfs/media/kasim/Data/data/FaceImage/replace1/FilterSelectFace.txt",
        help="path to image file list", type=str
    )
    parser.add_argument(
        "--merge_list", default="/rootfs/media/kasim/Data/data/FaceImage/filter/MergeFace.txt",
        help="path to merge faceid list", type=str
    )
    parser.add_argument(
        "--merge_image_list", default="/rootfs/media/kasim/Data/data/FaceImage/filter/MergeFaceImage.txt",
        help="path to merge face image file list", type=str
    )
    parser.add_argument(
        "--output_dir", default="/rootfs/media/kasim/Data/data/FaceImage/filter", help="path to mot reid image info",
        type=str
    )
    parser.add_argument('--file_list_mode', type=int, default=0, help='file list mode')
    parser.add_argument("--use_temp", action='store_true', help="use temp save, if crash can use it")
    parser.add_argument('--save_time', type=int, default=60, help='save time, second ')
    parser.add_argument('--max_display_count', type=int, default=8, help='max display count')
    args = parser.parse_args()

    random.seed(123456)

    save_time = args.save_time
    image_dir = args.image_dir
    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        os.system('chmod a+wr {}'.format(output_dir))

    display_image_count = h_image_count * v_image_count

    faceid_file_list_path = os.path.join(output_dir, 'select_merge_faces.txt')
    if not os.path.exists(faceid_file_list_path):
        xiaoqu_faceid_file_nums = read_merge_file(args.merge_list, args.merge_image_list, False)
    else:
        xiaoqu_faceid_file_nums = read_merge_file(faceid_file_list_path, args.merge_image_list)

    face_id_files = read_file_list(args.file_list, args.file_list_mode)

    cv2.namedWindow('image')
    cv2.setMouseCallback('image', mouse_event)

    last_id_idx = len(xiaoqu_faceid_file_nums) - 1

    cur_id_idx = 0
    cur_id_idx_file_path = os.path.join(output_dir, 'merge_cur_id_idx.txt')
    if os.path.exists(cur_id_idx_file_path):
        with open(cur_id_idx_file_path, 'r') as file:
            for line in file.readlines():
                cur_id_idx = int(line.strip())
                break

    if args.use_temp:
        tmp_output_dir = os.path.join('/dev/shm', 'SelectMergeImage')
        if not os.path.exists(tmp_output_dir):
            os.makedirs(tmp_output_dir)
            os.system('chmod a+wr {}'.format(tmp_output_dir))

        tmp_select_file_list_path = os.path.join(tmp_output_dir, 'select_merge_faces.txt')
        tmp_cur_id_idx_file_path = os.path.join(tmp_output_dir, 'merge_cur_id_idx.txt')

    display_more_images = False
    select_change = False
    start_save_time = time.time()
    last_face_count = None
    while cur_id_idx <= last_id_idx:
        xiaoqu, face_ids, file_list, face_nums = xiaoqu_faceid_file_nums[cur_id_idx]
        print(xiaoqu)
        last_idx = len(file_list) - 1
        start_idx = 0
        files = []

        if display_more_images:
            _face_count = len(face_ids)
            for i in range(_face_count):
                face_id = face_ids[i]
                face_num_files = face_id_files[face_id]
                dis_text = str(len(face_num_files)) + ':' + file_list[i]
                if len(face_num_files) > args.max_display_count:
                    face_num_files = random.sample(face_num_files, args.max_display_count)
                show_images(image_dir, face_num_files, dis_text=dis_text, win_name='{}'.format(i))
            if (last_face_count is not None) and (_face_count < last_face_count):
                for i in range(_face_count, last_face_count):
                    cv2.destroyWindow('{}'.format(i))
            last_face_count = _face_count

        for idx, file_name in enumerate(file_list):
            print('{}: {}'.format(idx, file_name))
            files.append(file_name)
            file_count = len(files)
            if file_count == 1:
                start_idx = idx
            if file_count < display_image_count and (idx < last_idx):
                continue
            imgs = []
            for file_name in files:
                file_path = os.path.join(image_dir, file_name)
                img = cv2.imread(file_path)
                img = image_resize(img, img_width, img_height)
                imgs.append(img)

            image_count = len(imgs)
            images = []
            x_off = 0
            y_off = 0
            image_bboxes.clear()
            for j, img in enumerate(imgs):
                if j % h_image_count == 0:
                    x_off = 0
                    if j != 0:
                        y_off += img_height
                    h_images = [img]
                else:
                    x_off += img_width
                    h_images.append(img)
                image_bboxes.append([x_off, y_off, (x_off + img_width), (y_off + img_height)])
                if ((j + 1) % h_image_count == 0) or ((j + 1) == image_count):
                    images.append(np.concatenate(h_images, axis=1))

            if len(images) > 1:
                w0 = images[0].shape[1]
                h0 = images[0].shape[0]
                w1 = images[-1].shape[1]
                h1 = images[-1].shape[0]
                if (w0 != w1) or (h0 != h1):
                    # image = np.empty_like(images[0])
                    image = np.zeros_like(images[0])
                    image[0:h1, 0:w1, :] = images[-1]
                    images[-1] = image
            draw_image = np.vstack(images)

            change_trk_id = False
            is_exit = False
            face_num_str = ''
            face_num = 1
            while True:
                image = draw_image.copy()
                cur_bbox_idx = get_xy_bbox_idx(current_mouse_xy)
                if cur_bbox_idx is None:
                    cur_bbox_idx = 0
                # print(cur_bbox_idx, current_mouse_xy)
                bbox = image_bboxes[cur_bbox_idx]
                cv2.rectangle(image, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 1)
                for select_idx in range(0, image_count):
                    bbox = image_bboxes[select_idx]
                    face_idx = start_idx + select_idx
                    num = face_nums[face_idx]
                    if 0 == num:
                        continue
                    cv2.putText(image, str(num), (bbox[0] + 2, bbox[3] - 4), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 255), thickness=2)
                cv2.putText(image, '%d/%d %d/%d' % (cur_id_idx + 1, last_id_idx + 1, idx + 1, last_idx + 1),
                            (8, 24), cv2.FONT_HERSHEY_PLAIN, 2.0, (0, 255, 0), thickness=2)
                cv2.imshow('image', image)
                key = cv2.waitKey(50)
                if key == 27:
                    is_exit = True
                    break
                elif key == ord(' '):  # 同一ID的下一页
                    select_change = True
                    break
                elif key == ord('a') or key == ord('A'):  # 上一个ID
                    select_change = True
                    if cur_id_idx > 0:
                        cur_id_idx -= 1
                        change_trk_id = True
                    break
                elif key == ord('d') or key == ord('D'):  # 下一个ID
                    select_change = True
                    change_trk_id = True
                    cur_id_idx += 1
                    break
                elif key == ord('c') or key == ord('C'):
                    cv2.destroyAllWindows()
                    cv2.namedWindow('image')
                    cv2.setMouseCallback('image', mouse_event)
                    continue
                elif ord('0') <= key <= ord('9'):  # 输入新FaceNUM
                    if face_num_str == '':
                        face_num_str = chr(key)
                    else:
                        face_num_str += chr(key)
                    continue
                elif key == ord('f') or key == ord('F'):  # 修改FaceNUM
                    if face_num_str != '':
                        face_num = int(face_num_str)
                        face_num_str = ''
                    select_idx = start_idx + cur_bbox_idx
                    face_nums[select_idx] = face_num
                    continue
                elif key == ord('g') or key == ord('G'):  # 所有人脸都改为同一个FaceNUM
                    if face_num_str != '':
                        face_num = int(face_num_str)
                        face_num_str = ''
                    face_nums_count = len(face_nums)
                    for select_idx in range(face_nums_count):
                        face_nums[select_idx] = face_num
                    continue
                elif key == ord('r') or key == ord('r'):  # 修改FaceNUM+1
                    face_num += 1
                    select_idx = start_idx + cur_bbox_idx
                    face_nums[select_idx] = face_num
                    continue
                elif key == ord('e') or key == ord('E'):  # 删除选择人脸
                    select_idx = start_idx + cur_bbox_idx
                    face_nums[select_idx] = -1
                    continue
                elif key == ord('t') or key == ord('T'):  # 删除所有人脸
                    face_nums_count = len(face_nums)
                    for select_idx in range(face_nums_count):
                        face_nums[select_idx] = -1
                    continue
                elif key == ord('v') or key == ord('V'):  # 此人脸下的图片不全为一个人，需要重新标注
                    select_idx = start_idx + cur_bbox_idx
                    face_nums[select_idx] = -2
                    continue
                elif key == ord('w') or key == ord('W'):  # 不好区分此人脸与其他人脸是否一样
                    select_idx = start_idx + cur_bbox_idx
                    face_nums[select_idx] = -3
                    continue
                elif key == ord('x') or key == ord('X'):  # 显示选中的人脸对应的face_num所对应的所有人脸的所有图片
                    select_idx = start_idx + cur_bbox_idx
                    _face_num = face_nums[select_idx]
                    cv2.destroyAllWindows()
                    cv2.namedWindow('image')
                    cv2.setMouseCallback('image', mouse_event)
                    for i, num in enumerate(face_nums):
                        if num == _face_num:
                            face_id = face_ids[i]
                            face_num_files = face_id_files[face_id]
                            show_images(image_dir, face_num_files, win_name='{}'.format(i))
                    continue
                elif key == ord('q') or key == ord('Q'):  # 每个人显示多大max_display_count张图片
                    display_more_images = not display_more_images
                    if display_more_images:
                        _face_count = len(face_ids)
                        for i in range(_face_count):
                            face_id = face_ids[i]
                            face_num_files = face_id_files[face_id]
                            dis_text = str(len(face_num_files)) + ':' + file_list[i]
                            if len(face_num_files) > args.max_display_count:
                                face_num_files = random.sample(face_num_files, args.max_display_count)
                            show_images(image_dir, face_num_files, dis_text=dis_text, win_name='{}'.format(i))
                        last_face_count = _face_count
                    else:
                        cv2.destroyAllWindows()
                        cv2.namedWindow('image')
                        cv2.setMouseCallback('image', mouse_event)
                    continue
                elif key == ord('z') or key == ord('Z'):  # 显示选中的人脸的所有图片
                    select_idx = start_idx + cur_bbox_idx
                    # cv2.destroyAllWindows()
                    # cv2.namedWindow('image')
                    # cv2.setMouseCallback('image', mouse_event)
                    face_id = face_ids[select_idx]
                    face_num_files = face_id_files[face_id]
                    show_images(image_dir, face_num_files, win_name='{}'.format(select_idx))
                    continue
                elif key == ord('s') or key == ord('S'):  # Save
                    if args.use_temp:
                        save_process(tmp_select_file_list_path, xiaoqu_faceid_file_nums, tmp_cur_id_idx_file_path,
                                     cur_id_idx)
                    start_save_time = time.time()
                    select_change = False
                    break
                if args.use_temp:
                    # 定时保存
                    if time.time() - start_save_time >= save_time:
                        if select_change:
                            save_process(tmp_select_file_list_path, xiaoqu_faceid_file_nums, tmp_cur_id_idx_file_path,
                                         cur_id_idx)
                            select_change = False
                        start_save_time = time.time()

            files.clear()
            if is_exit:
                save_process(faceid_file_list_path, xiaoqu_faceid_file_nums, cur_id_idx_file_path, cur_id_idx)
                return
            if change_trk_id:
                break
    save_process(faceid_file_list_path, xiaoqu_faceid_file_nums, cur_id_idx_file_path, cur_id_idx)
    print('finish!')


if __name__ == '__main__':
    main()
